import mne
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import spectrogram
import os

# 设置 Matplotlib 支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def process_and_plot_ssvep(
        vhdr_file_path: str,
        channels_to_average: list[str],
        output_dir: str,
        output_filename: str,
        analysis_duration: float = 16.0,
        # 每个时间窗口的长度(小-时间刻度（X轴）会更细)
        nperseg_sec: float = 2.0,
        # 相邻窗口有 秒的重叠
        noverlap_sec: float = 0.75,
        # 使用前2秒作为基线
        baseline_sec: tuple[float, float] = (0.0, 1.0)
):
    """
    加载 SSVEP 数据，裁剪到事件开始后的指定时间段，
    对指定通道进行平均，计算时频能量图并保存。

    参数:
    vhdr_file_path (str): BrainVision .vhdr 文件的完整路径。
    channels_to_average (list[str]): 要进行平均并绘制时频图的通道名称列表。
    output_dir (str): 保存输出图的文件夹路径。
    output_filename (str): 输出图的文件名。
    analysis_duration (float): 从事件开始后分析的总时长（秒）。
    nperseg_sec (float): STFT 窗口长度（秒）。
    noverlap_sec (float): STFT 窗口重叠时长（秒）。
    baseline_sec (tuple[float, float]): 用于基线校正的时间窗（秒）。
    """
    print(f"开始分析文件: {vhdr_file_path}")

    # --- 1. 数据加载 ---
    try:
        raw = mne.io.read_raw_brainvision(vhdr_file_path, preload=True)
        print("数据加载成功。")
    except FileNotFoundError:
        print(f"错误：未找到文件 '{vhdr_file_path}'。请确保所有相关文件 (.vhdr, .eeg, .vmrk) 在同一目录中。")
        return

    # 设置标准电极位置并忽略未找到的通道
    montage = mne.channels.make_standard_montage('standard_1020')
    raw.set_montage(montage, on_missing='ignore')

    # --- 2. 查找事件开始时间并裁剪数据 ---
    print("正在查找事件开始时间并裁剪数据...")
    events, _ = mne.events_from_annotations(raw)

    if len(events) == 0:
        print("警告：未找到任何事件。将分析整个数据记录。")
        start_time = 0
    else:
        start_time = events[0, 0] / raw.info['sfreq']
        print(f"事件开始时间为：{start_time:.2f} 秒。")

    end_time = start_time + analysis_duration

    raw_cropped = raw.copy().crop(tmin=start_time, tmax=end_time)
    print(f"数据已裁剪至 {start_time:.2f} 秒到 {end_time:.2f} 秒。")

    valid_channels = [ch for ch in channels_to_average if ch in raw_cropped.ch_names]
    if not valid_channels:
        print(f"错误：数据中未找到任何指定的通道：{channels_to_average}。无法继续分析。")
        return

    # --- 3. 预处理 ---
    print("开始数据预处理...")
    raw_cropped.filter(1., 40., fir_design='firwin')

    eeg_data = raw_cropped.copy().pick(valid_channels)
    data_matrix = eeg_data.get_data()
    sampling_freq = raw_cropped.info['sfreq']
    print("预处理完成。")

    # --- 4. 通道平均和时频分析 ---
    print("开始通道平均和时频分析...")
    averaged_data = np.mean(data_matrix, axis=0)

    nperseg = int(nperseg_sec * sampling_freq)
    noverlap = int(noverlap_sec * sampling_freq)

    if len(averaged_data) < nperseg:
        print(f"警告: 数据长度 ({len(averaged_data)}) 小于STFT窗口长度 ({nperseg})。跳过绘图。")
        return

    f, t, Sxx = spectrogram(averaged_data, sampling_freq, nperseg=nperseg, noverlap=noverlap)
    print("时频分析完成。")
    print(f)

    # --- 5. 基线校正 ---
    print("正在进行基线校正...")
    t_start, t_end = baseline_sec
    baseline_idx = (t >= t_start) & (t <= t_end)

    if np.sum(baseline_idx) == 0:
        print("警告: 找不到有效的基线时间窗。跳过基线校正。")
        Sxx_db = 10 * np.log10(Sxx)
    else:
        # 计算基线期内每个频率的平均能量
        baseline_power_mean = np.mean(Sxx[:, baseline_idx], axis=1, keepdims=True)

        # 将能量转换为 dB，并减去基线平均能量
        Sxx_db = 10 * np.log10(Sxx / baseline_power_mean)
        print("基线校正完成。")

    # --- 6. 绘图和保存 ---
    print("生成并保存时频图...")

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"已创建目录：'{output_dir}'")

    plt.figure(figsize=(12, 8))
    # 设置能量的动态范围
    vmax = Sxx_db.max()
    vmin = Sxx_db.min()
    plt.pcolormesh(t, f, Sxx_db, shading='gouraud', cmap='viridis', vmin=vmin, vmax=vmax)

    plt.title(
        f'SSVEP 时频能量图 - 通道: {valid_channels}\n'
        f'时间范围: {start_time:.2f}-{end_time:.2f} 秒 | 基线校正: {t_start:.1f}-{t_end:.1f} 秒'
    )
    plt.ylabel('频率 [Hz]')
    plt.xlabel('时间 [秒]')
    plt.ylim([5, 25])
    plt.colorbar(label='能量变化 (dB)')

    output_path = os.path.join(output_dir, output_filename)
    plt.savefig(output_path, dpi=300)

    print(f"成功生成并保存图表至 '{output_path}'")
    plt.show()


# --- 主程序执行部分 ---
if __name__ == '__main__':
    vhdr_file = '../data/hzy/hzy_ssvep_2.vhdr'
    channels = ['O2','Oz','O1','P7','P3','Pz','P4','P8']
    output_folder = 'pic/gz_2'
    plot_name = 'gz_ssvep_4.png'

    process_and_plot_ssvep(
        vhdr_file_path=vhdr_file,
        channels_to_average=channels,
        output_dir=output_folder,
        output_filename=plot_name
    )